import numpy as np
from scipy.signal import iirnotch, lfilter
import matplotlib.pyplot as plt

# 设计陷波滤波器
def design_notch_filter(freq, fs, quality):
    """
    设计一个陷波滤波器
    :param freq: 要抑制的频率 (Hz)，例如 50Hz
    :param fs: 采样率 (Hz)
    :param quality: 品质因数 (Q)，影响滤波器的带宽
    :return: 滤波器的系数 b, a
    """
    b, a = iirnotch(w0=freq / (fs / 2), Q=quality)
    return b, a

# 模拟信号
fs = 1000  # 采样率 (Hz)
t = np.linspace(0, 1, fs, endpoint=False)
signal = np.sin(2 * np.pi * 50 * t) + 0.5 * np.sin(2 * np.pi * 120 * t)  # 含50Hz噪声的信号

# 应用陷波滤波器
import time
mytime = time.perf_counter()
notch_freq = 50  # 工频
quality_factor = 30  # 品质因数
b, a = design_notch_filter(notch_freq, fs, quality_factor)
filtered_signal = lfilter(b, a, signal)
curtime = time.perf_counter()
print(f"total {curtime - mytime}s")

# 可视化
plt.figure(figsize=(10, 6))
plt.subplot(2, 1, 1)
plt.title("before")
plt.plot(t, signal)
plt.subplot(2, 1, 2)
plt.title("after")
plt.plot(t, filtered_signal)
plt.tight_layout()
plt.show()
